/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lxq.cloud.ai.controller;


import com.lxq.cloud.ai.service.ImagesService;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.ai.image.*;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.nio.charset.StandardCharsets;

@RestController
@RequestMapping("/ai")
public class ImageModelController {
    private final ImagesService imagesService;
    private final ImageModel imageModel;

    ImageModelController(ImageModel imageModel, ImagesService imagesService) {
        this.imageModel = imageModel;
        this.imagesService = imagesService;
    }

    @GetMapping("/image/{input}")
    public void image(@PathVariable("input") String input, HttpServletResponse response) {

        ImageOptions options = ImageOptionsBuilder.builder().model("wanx-v1").build();

        ImagePrompt imagePrompt = new ImagePrompt(input, options);
        ImageResponse imageResponse = imageModel.call(imagePrompt);
        String imageUrl = imageResponse.getResult().getOutput().getUrl();

        try {
            URL url = new URL(imageUrl);
            InputStream in = url.openStream();

            response.setHeader("Content-Type", MediaType.IMAGE_PNG_VALUE);
            response.getOutputStream().write(in.readAllBytes());
            response.getOutputStream().flush();
        } catch (IOException e) {
            response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
        }
    }


    @GetMapping("/img")
    public ResponseEntity<byte[]> genImg(@RequestParam(value = "prompt", defaultValue = "Painting a picture of blue water and blue sky.") String imgPrompt) {
        try {
            ImageResponse imageResponse = imagesService.genImg(imgPrompt);

            URL url = new URL(imageResponse.getResult().getOutput().getUrl());
            try (InputStream in = url.openStream()) {
                byte[] imageBytes = in.readAllBytes();
                // 设置 Content-Type（根据图像格式） 若是 PNG 图片
                final MediaType mediaType = MediaType.IMAGE_PNG;

                return ResponseEntity.ok()
                        .contentType(mediaType)
                        .body(imageBytes);
            }
        } catch (final Exception e) {
            // 处理异常，返回错误响应
            return ResponseEntity.status(500)
                    .contentType(MediaType.TEXT_PLAIN)
                    .body("图像生成失败: ".getBytes(StandardCharsets.UTF_8));
        }
    }
}
